Skip to content

Neighborhood filter kernels caching/race condition fixes - #1768

Open
cmdupuis3 wants to merge 12 commits into
mainfrom
cmd/module-level-neighborhood-kernels
Open

cmdupuis3 wants to merge 12 commits into
mainfrom
cmd/module-level-neighborhood-kernels

Conversation

@cmdupuis3

@cmdupuis3 cmdupuis3 commented Sep 16, 2026 •

Copy link
Copy Markdown
Collaborator

Closes #1757, #1769

Overview

The PR solves two issues with the current version of the neighborhood filters kernels:

  1. numba has issues caching closures that are also decorated as njit(cache=True). This means that in the current implementation of kernel caching, every call to a kernel gets a new UUID, so any attempt to cache this way will basically be a no-op.
  2. Somewhat related, but there is also a race condition with dask. See Race condition crashes pytest suite locally sometimes #1757.

In a fun twist, both are known unresolved issues in numba:

numba/numba#6264
numba/numba#9288

PR Checklist

General

  • An issue is created and linked
  • Added appropriate labels (if your uxarray repo permissions allow it)
  • Filled out Overview and Expected Usage (if applicable) sections

Testing & Benchmarking

  • There is adequate test coverage of changes from this PR (add new tests if needed)
  • If this PR could affect performance, ran ASV benchmarks and confirmed they show expected behavior (add a new benchmark if necessary)

Documentation and Examples

  • Docstrings updated with any function changes, and included in all new functions
  • User (public) functions added to docs/api.rst; internal (private) function names start with an underscore (_)

AI Disclosure

AI Usage: Claude Opus 5

  • I have tested and take responsibility for all AI-generated content in my PR.

cmdupuis3 and others added 3 commits September 14, 2026 13:41
The neighborhood gufuncs are built with ``cache=True``, but the cache never
hit: every process recompiled all nine kernels and appended a fresh set of
entries to the index, which grew without bound (over a thousand files in a
working checkout).

The cause is that ``_make_kernel`` closed each kernel over its reducer. Numba
keys a cached function on a hash of its closure, and a ``Dispatcher``
serializes with a ``uuid4`` regenerated in every process, so the key differed
in every process. Numba reports no error for this -- it writes the cache and
silently misses it. All nine kernels also shared one cache file, since the
closure made them one function as far as numba's locator was concerned.

Give each reduction a module-level kernel body instead, and apply
``guvectorize`` to it on first call. A module-level body has no closure, so its
key is stable, and it reaches its reducer as a global, which numba resolves at
compile time and leaves out of the key. Each body also gets its own cache file.
Kernels are now compiled once per machine rather than once per process.

Compilation stays lazy. ``guvectorize`` compiles at decoration time when given
explicit signatures, so decorating at module scope would rebuild every kernel
during ``import uxarray`` and start numba's threading layer, leaving a thread
pool that makes forks unsafe -- the regression
``test_no_numba_kernels_built_on_import`` guards against.

Replace ``functools.cache`` with a double-checked lock while here. It does not
hold a lock across the call it memoizes, so under ``dask="parallelized"`` every
worker thread that reached a kernel before the first build finished started its
own full compilation, serialized behind numba's global compiler lock. A 12-chunk
reduction compiled the same kernel 12 times; it now compiles once.

The bodies are spelled out rather than generated because the deduplication one
would reach for -- a single shared body taking the reducer as an argument --
makes the reducer a dynamic global, which numba refuses to cache at all. The
gather is shared through ``_widest``/``_gather`` instead, leaving only the
reducer name different between bodies.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Replaces the ``_lazy`` closure with a ``_LazyKernel`` descriptor that holds the
compiled gufunc on itself, so each kernel is still built on first use but the
memoization lives at class level rather than in a per-kernel closure.

Two things fall out. The ``staticmethod`` wrapping goes away: a descriptor
hands back the gufunc itself, so ``self`` is never bound as the kernel's first
argument, and the comment explaining that wart goes with it. And compilation
now happens where the attribute is resolved -- ``_apply_kernel`` reads it on
the calling thread -- so a dask-backed reduction builds its kernel while the
graph is being assembled rather than inside a task. The array stays lazy; only
the compile moves, and it moves out of the parallel region, which makes the
worker-thread race structurally impossible rather than merely locked against.

``functools.cached_property`` would not do. It caches per instance, and
``Grid.neighborhood()`` returns a new ``Neighborhood`` on every call, so it ran
``guvectorize`` once per neighborhood: fifty fresh instances measured fifty
compilations, against zero for the descriptor. It also holds no lock, CPython
having removed it in 3.12, so twelve threads racing one instance measured
twelve compilations, against one.

The kernel bodies must still be module-level functions with no closure. That is
what keeps numba's cache key stable across processes, and it is independent of
where the memoization lives -- the memoizer never enters the key.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
@cmdupuis3 cmdupuis3 self-assigned this Sep 16, 2026
@cmdupuis3
cmdupuis3 marked this pull request as draft September 16, 2026 22:37
@cmdupuis3 cmdupuis3 linked an issue Sep 16, 2026 that may be closed by this pull request
@cmdupuis3 cmdupuis3 added bug Something isn't working scalability Related to scalability & performance efforts run-benchmark Run ASV benchmark workflow benchmarking Related to benchmarks, memory usage, and/or time profiling labels Sep 16, 2026
@cmdupuis3
cmdupuis3 requested a review from Sevans711 September 16, 2026 22:46
@github-actions

github-actions Bot commented Sep 16, 2026 •

Copy link
Copy Markdown

ASV Benchmarking

Benchmark Comparison Results

Benchmarks that have improved:

Change Before [e01d71e] After [c2b3706] Ratio Benchmark (Parameter)
- 348M 275M 0.79 import.Imports.track_peakmem_import_uxarray
- 348M 312M 0.9 mpas_ocean.GradientColdStartRss.track_peakmem_gradient('480km')
- 44.4±1ms 40.3±0.5ms 0.91 mpas_ocean.NeighborhoodDask.time_mean('120km', 'grid_chunks')
- 22.7±0.02ms 18.4±0.1ms 0.81 mpas_ocean.NeighborhoodDask.time_mean('120km', 'numpy')
- 38.5±0.2ms 31.0±0.1ms 0.81 mpas_ocean.NeighborhoodReduce.time_reduce('120km', 'mean')

Benchmarks that have stayed the same:

Change Before [e01d71e] After [c2b3706] Ratio Benchmark (Parameter)
5.30±0.02ms 5.32±0.04ms 1.00 bench_connectivity.Connectivity.time_edge_face('120km')
2.01±0.01ms 1.99±0.1ms 0.99 bench_connectivity.Connectivity.time_edge_face('480km')
4.48±0.08ms 4.40±0.04ms 0.98 bench_connectivity.Connectivity.time_edge_node('120km')
1.58±0.02ms 1.58±0.02ms 1.00 bench_connectivity.Connectivity.time_edge_node('480km')
4.50±0.09ms 4.46±0.1ms 0.99 bench_connectivity.Connectivity.time_face_edge('120km')
1.60±0.02ms 1.57±0.01ms 0.98 bench_connectivity.Connectivity.time_face_edge('480km')
6.26±0.1ms 6.26±0.04ms 1.00 bench_connectivity.Connectivity.time_face_face('120km')
2.34±0ms 2.38±0.01ms 1.02 bench_connectivity.Connectivity.time_face_face('480km')
54.0±2μs 51.4±0.6μs 0.95 bench_connectivity.Connectivity.time_face_node('120km')
53.3±2μs 50.9±1μs 0.95 bench_connectivity.Connectivity.time_face_node('480km')
449±10μs 425±20μs 0.95 bench_connectivity.Connectivity.time_n_nodes_per_face('120km')
369±10μs 362±10μs 0.98 bench_connectivity.Connectivity.time_n_nodes_per_face('480km')
5.85±0.1ms 5.72±0.08ms 0.98 bench_connectivity.Connectivity.time_node_edge('120km')
2.01±0.01ms 1.98±0ms 0.99 bench_connectivity.Connectivity.time_node_edge('480km')
81.2±4ms 79.9±1ms 0.98 bench_connectivity.Connectivity.time_node_face('120km')
5.20±0.05ms 5.12±0.05ms 0.99 bench_connectivity.Connectivity.time_node_face('480km')
8.44±0.2ms 8.43±0.2ms 1.00 face_bounds.FaceBounds.time_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/mpas/QU/oQU480.231010.nc'))
2.74±0.04ms 2.78±0.06ms 1.01 face_bounds.FaceBounds.time_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/scrip/outCSne8/outCSne8.nc'))
7.12±7s 10.3±10ms ~0.00 face_bounds.FaceBounds.time_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/geoflow-small/grid.nc'))
1.60±0.04ms 1.55±0.04ms 0.97 face_bounds.FaceBounds.time_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/quad-hexagon/grid.nc'))
57.3k 57.3k 1.00 face_bounds.FaceBounds.track_nbytes_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/mpas/QU/oQU480.231010.nc'))
12.3k 12.3k 1.00 face_bounds.FaceBounds.track_nbytes_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/scrip/outCSne8/outCSne8.nc'))
123k 123k 1.00 face_bounds.FaceBounds.track_nbytes_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/geoflow-small/grid.nc'))
128 128 1.00 face_bounds.FaceBounds.track_nbytes_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/quad-hexagon/grid.nc'))
1.27M 1.27M 1.00 face_bounds.FaceBounds.track_nbytes_grid_with_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/mpas/QU/oQU480.231010.nc'))
50.1k 50.1k 1.00 face_bounds.FaceBounds.track_nbytes_grid_with_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/scrip/outCSne8/outCSne8.nc'))
1.48M 1.48M 1.00 face_bounds.FaceBounds.track_nbytes_grid_with_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/geoflow-small/grid.nc'))
712 712 1.00 face_bounds.FaceBounds.track_nbytes_grid_with_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/quad-hexagon/grid.nc'))
1.98M 1.98M 1.00 face_bounds.FaceBounds.track_peakmem_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/mpas/QU/oQU480.231010.nc'))
1.97M 1.97M 1.00 face_bounds.FaceBounds.track_peakmem_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/scrip/outCSne8/outCSne8.nc'))
2.13M 2.13M 1.00 face_bounds.FaceBounds.track_peakmem_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/geoflow-small/grid.nc'))
35.5k 35.5k 1.00 face_bounds.FaceBounds.track_peakmem_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/quad-hexagon/grid.nc'))
348M 317M 0.91 face_bounds.FaceBoundsColdStartRss.track_peakmem_open_and_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/mpas/QU/oQU480.231010.nc'))
363M 361M 0.99 face_bounds.FaceBoundsColdStartRss.track_peakmem_open_and_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/scrip/outCSne8/outCSne8.nc'))
348M 319M 0.92 face_bounds.FaceBoundsColdStartRss.track_peakmem_open_and_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/geoflow-small/grid.nc'))
348M 318M 0.91 face_bounds.FaceBoundsColdStartRss.track_peakmem_open_and_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/quad-hexagon/grid.nc'))
1.17±0.02μs 1.15±0.05μs 0.98 geometry_kernels.AccucrossKernels.time_accucross
2.59±0.02μs 2.57±0.03μs 0.99 geometry_kernels.AccucrossKernels.time_accucross_pair
465±20ns 436±20ns 0.94 geometry_kernels.EFTPrimitives.time_acc_sqrt_re
451±20ns 436±10ns 0.97 geometry_kernels.EFTPrimitives.time_diff_of_products
391±30ns 386±20ns 0.99 geometry_kernels.EFTPrimitives.time_two_prod
396±30ns 376±20ns 0.95 geometry_kernels.EFTPrimitives.time_two_sum
691±10ns 677±20ns 0.98 geometry_kernels.GCAConstLatIntersection.time_accux_constlat_kernel
766±20ns 726±20ns 0.95 geometry_kernels.GCAConstLatIntersection.time_gca_const_lat_intersection
847±20ns 792±10ns 0.94 geometry_kernels.GCAConstLatIntersection.time_try_gca_const_lat_intersection
846±10ns 801±20ns 0.95 geometry_kernels.GCAGCAIntersection.time_accux_gca_kernel
932±10ns 912±5ns 0.98 geometry_kernels.GCAGCAIntersection.time_gca_gca_intersection
1.06±0.02μs 1.06±0.03μs 1.00 geometry_kernels.GCAGCAIntersection.time_try_gca_gca_intersection
53.5±1μs 50.3±0.5μs 0.94 geometry_kernels.OrientPredicates.time_on_minor_arc
52.8±2μs 49.6±0.4μs 0.94 geometry_kernels.OrientPredicates.time_orient3d_on_sphere
3.07±0ms 3.07±0.01ms 1.00 geometry_samebody.SameBodyConstLat.time_accux_dispatch
1.16±0ms 1.16±0ms 1.00 geometry_samebody.SameBodyConstLat.time_accux_kernel
2.30±0.01ms 2.29±0.01ms 1.00 geometry_samebody.SameBodyConstLat.time_fp64_dispatch
147±1μs 146±0.07μs 1.00 geometry_samebody.SameBodyConstLat.time_fp64_kernel
29.6±0.04ms 30.0±0.04ms 1.01 geometry_samebody_gcagca.SameBodyGcaGca.time_accux_dispatch
6.32±0.01ms 6.28±0ms 1.00 geometry_samebody_gcagca.SameBodyGcaGca.time_accux_kernel
23.7±0.09ms 23.6±0.04ms 1.00 geometry_samebody_gcagca.SameBodyGcaGca.time_fp64_dispatch
858±0.7μs 861±1μs 1.00 geometry_samebody_gcagca.SameBodyGcaGca.time_fp64_kernel
906±4ms 915±6ms 1.01 import.Imports.timeraw_import_uxarray
2.55±0.02ms 2.50±0.05ms 0.98 mpas_ocean.CheckNorm.time_check_norm('120km')
2.13±0.07ms 2.00±0.03ms 0.94 mpas_ocean.CheckNorm.time_check_norm('480km')
1.15±0.01ms 1.15±0.01ms 1.00 mpas_ocean.ConnectivityConstruction.time_face_face_connectivity('120km')
564±4μs 569±10μs 1.01 mpas_ocean.ConnectivityConstruction.time_face_face_connectivity('480km')
657±9μs 648±9μs 0.98 mpas_ocean.ConnectivityConstruction.time_n_nodes_per_face('120km')
606±9μs 604±7μs 1.00 mpas_ocean.ConnectivityConstruction.time_n_nodes_per_face('480km')
5.31±0.03ms 5.61±0.06ms 1.06 mpas_ocean.ConstructFaceLatLon.time_cartesian_averaging('120km')
3.86±0.02ms 3.92±0.2ms 1.01 mpas_ocean.ConstructFaceLatLon.time_cartesian_averaging('480km')
97.0±0.3ms 97.2±0.4ms 1.00 mpas_ocean.ConstructFaceLatLon.time_welzl('120km')
10.1±0.2ms 9.83±0.5ms 0.97 mpas_ocean.ConstructFaceLatLon.time_welzl('480km')
18.8±0.1ms 18.9±0.1ms 1.00 mpas_ocean.ConstructTreeStructures.time_ball_tree('120km')
1.00±0.03ms 997±40μs 0.99 mpas_ocean.ConstructTreeStructures.time_ball_tree('480km')
10.1±0.04ms 9.99±0.09ms 0.99 mpas_ocean.ConstructTreeStructures.time_kd_tree('120km')
634±20μs 614±20μs 0.97 mpas_ocean.ConstructTreeStructures.time_kd_tree('480km')
616±10ms 600±10ms 0.97 mpas_ocean.CrossSections.time_const_lat('120km', 1)
304±0.5ms 306±2ms 1.00 mpas_ocean.CrossSections.time_const_lat('120km', 2)
154±0.7ms 154±2ms 1.00 mpas_ocean.CrossSections.time_const_lat('120km', 4)
552±2ms 542±10ms 0.98 mpas_ocean.CrossSections.time_const_lat('480km', 1)
286±0.6ms 275±4ms 0.96 mpas_ocean.CrossSections.time_const_lat('480km', 2)
142±2ms 141±1ms 1.00 mpas_ocean.CrossSections.time_const_lat('480km', 4)
348M 337M 0.97 mpas_ocean.CrossSectionsPeakMem.track_peakmem_const_lat('120km', 1)
348M 337M 0.97 mpas_ocean.CrossSectionsPeakMem.track_peakmem_const_lat('120km', 2)
348M 337M 0.97 mpas_ocean.CrossSectionsPeakMem.track_peakmem_const_lat('120km', 4)
348M 320M 0.92 mpas_ocean.CrossSectionsPeakMem.track_peakmem_const_lat('480km', 1)
348M 320M 0.92 mpas_ocean.CrossSectionsPeakMem.track_peakmem_const_lat('480km', 2)
348M 320M 0.92 mpas_ocean.CrossSectionsPeakMem.track_peakmem_const_lat('480km', 4)
26.0±0.3ms 26.2±0.2ms 1.01 mpas_ocean.DualMesh.time_dual_mesh_construction('120km')
3.00±0.05ms 2.95±0.05ms 0.98 mpas_ocean.DualMesh.time_dual_mesh_construction('480km')
14.5±0.4ms 14.9±0.8ms 1.02 mpas_ocean.FaceAreas.time_face_areas('120km')
4.32±0.4ms 4.46±0.4ms 1.03 mpas_ocean.FaceAreas.time_face_areas('480km')
229k 229k 1.00 mpas_ocean.FaceAreas.track_nbytes_face_areas('120km')
14.3k 14.3k 1.00 mpas_ocean.FaceAreas.track_nbytes_face_areas('480km')
2.12M 2.12M 1.00 mpas_ocean.FaceAreas.track_peakmem_face_areas('120km')
714k 715k 1.00 mpas_ocean.FaceAreas.track_peakmem_face_areas('480km')
912±9ms 908±8ms 0.99 mpas_ocean.GeoDataFrame.time_to_geodataframe('120km', False)
52.1±1ms 52.2±0.3ms 1.00 mpas_ocean.GeoDataFrame.time_to_geodataframe('120km', True)
82.4±2ms 82.9±2ms 1.01 mpas_ocean.GeoDataFrame.time_to_geodataframe('480km', False)
5.02±0.3ms 5.30±0.1ms 1.05 mpas_ocean.GeoDataFrame.time_to_geodataframe('480km', True)
13.2±0.2ms 13.9±0.5ms 1.05 mpas_ocean.Gradient.time_gradient('120km')
1.81±0.03ms 1.79±0.01ms 0.99 mpas_ocean.Gradient.time_gradient('480km')
457k 457k 1.00 mpas_ocean.Gradient.track_nbytes_gradient('120km')
28.7k 28.7k 1.00 mpas_ocean.Gradient.track_nbytes_gradient('480km')
3.2M 3.2M 1.00 mpas_ocean.Gradient.track_peakmem_gradient('120km')
204k 204k 1.00 mpas_ocean.Gradient.track_peakmem_gradient('480km')
348M 335M 0.96 mpas_ocean.GradientColdStartRss.track_peakmem_gradient('120km')
276±8μs 289±20μs 1.05 mpas_ocean.HoleEdgeIndices.time_construct_hole_edge_indices('120km')
148±1μs 146±1μs 0.99 mpas_ocean.HoleEdgeIndices.time_construct_hole_edge_indices('480km')
200±2μs 200±1μs 1.00 mpas_ocean.Integrate.time_integrate('120km')
186±3μs 187±10μs 1.01 mpas_ocean.Integrate.time_integrate('480km')
18.4M 18.4M 1.00 mpas_ocean.Integrate.track_nbytes_integrate('120km')
1.2M 1.2M 1.00 mpas_ocean.Integrate.track_nbytes_integrate('480km')
187±2ms 189±1ms 1.01 mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('120km', 'exclude')
184±1ms 185±0.8ms 1.00 mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('120km', 'include')
187±1ms 188±1ms 1.01 mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('120km', 'split')
13.1±0.06ms 13.2±0.2ms 1.01 mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('480km', 'exclude')
13.3±0.2ms 13.3±0.4ms 1.00 mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('480km', 'include')
13.3±0.08ms 13.3±0.08ms 1.00 mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('480km', 'split')
247±1ms 247±1ms 1.00 mpas_ocean.NeighborhoodBuild.time_build('120km', 1.0)
1.31±0.01s 1.32±0.01s 1.01 mpas_ocean.NeighborhoodBuild.time_build('120km', 15.0)
509±4ms 507±2ms 1.00 mpas_ocean.NeighborhoodBuild.time_build('120km', 5.0)
13.5±0.1ms 13.5±0.1ms 1.00 mpas_ocean.NeighborhoodBuild.time_build('480km', 1.0)
25.2±0.1ms 25.4±0.3ms 1.01 mpas_ocean.NeighborhoodBuild.time_build('480km', 15.0)
16.7±0.2ms 16.6±0.3ms 1.00 mpas_ocean.NeighborhoodBuild.time_build('480km', 5.0)
241±2ms 243±1ms 1.01 mpas_ocean.NeighborhoodBuild.time_query_radius('120km', 1.0)
1.29±0s 1.28±0.01s 0.99 mpas_ocean.NeighborhoodBuild.time_query_radius('120km', 15.0)
499±3ms 504±3ms 1.01 mpas_ocean.NeighborhoodBuild.time_query_radius('120km', 5.0)
13.0±0.06ms 13.0±0.09ms 1.00 mpas_ocean.NeighborhoodBuild.time_query_radius('480km', 1.0)
25.1±0.03ms 25.0±0.1ms 1.00 mpas_ocean.NeighborhoodBuild.time_query_radius('480km', 15.0)
16.2±0.08ms 16.1±0.1ms 0.99 mpas_ocean.NeighborhoodBuild.time_query_radius('480km', 5.0)
1.19 1.19 1.00 mpas_ocean.NeighborhoodBuild.track_mean_neighbors('120km', 1.0)
612.76 612.76 1.00 mpas_ocean.NeighborhoodBuild.track_mean_neighbors('120km', 15.0)
74.17 74.17 1.00 mpas_ocean.NeighborhoodBuild.track_mean_neighbors('120km', 5.0)
1.0 1.0 1.00 mpas_ocean.NeighborhoodBuild.track_mean_neighbors('480km', 1.0)
37.29 37.29 1.00 mpas_ocean.NeighborhoodBuild.track_mean_neighbors('480km', 15.0)
6.57 6.57 1.00 mpas_ocean.NeighborhoodBuild.track_mean_neighbors('480km', 5.0)
728k 728k 1.00 mpas_ocean.NeighborhoodBuild.track_nbytes_neighbors('120km', 1.0)
141M 141M 1.00 mpas_ocean.NeighborhoodBuild.track_nbytes_neighbors('120km', 15.0)
17.4M 17.4M 1.00 mpas_ocean.NeighborhoodBuild.track_nbytes_neighbors('120km', 5.0)
43k 43k 1.00 mpas_ocean.NeighborhoodBuild.track_nbytes_neighbors('480km', 1.0)
563k 563k 1.00 mpas_ocean.NeighborhoodBuild.track_nbytes_neighbors('480km', 15.0)
123k 123k 1.00 mpas_ocean.NeighborhoodBuild.track_nbytes_neighbors('480km', 5.0)
5.72M 5.72M 1.00 mpas_ocean.NeighborhoodBuild.track_peakmem_build('120km', 1.0)
145M 145M 1.00 mpas_ocean.NeighborhoodBuild.track_peakmem_build('120km', 15.0)
21.5M 21.5M 1.00 mpas_ocean.NeighborhoodBuild.track_peakmem_build('120km', 5.0)
362k 362k 1.00 mpas_ocean.NeighborhoodBuild.track_peakmem_build('480km', 1.0)
825k 825k 1.00 mpas_ocean.NeighborhoodBuild.track_peakmem_build('480km', 15.0)
384k 384k 1.00 mpas_ocean.NeighborhoodBuild.track_peakmem_build('480km', 5.0)
40.6±0.3ms 36.9±1ms 0.91 mpas_ocean.NeighborhoodDask.time_mean('120km', 'time_chunks')
11.9±0.1ms 11.2±0.2ms 0.93 mpas_ocean.NeighborhoodDask.time_mean('480km', 'grid_chunks')
8.39±0.2ms 7.72±0.1ms 0.92 mpas_ocean.NeighborhoodDask.time_mean('480km', 'time_chunks')
5.84M 5.74M 0.98 mpas_ocean.NeighborhoodDask.track_peakmem_mean('120km', 'grid_chunks')
2.75M 2.75M 1.00 mpas_ocean.NeighborhoodDask.track_peakmem_mean('120km', 'numpy')
5.69M 5.68M 1.00 mpas_ocean.NeighborhoodDask.track_peakmem_mean('120km', 'time_chunks')
177k 177k 1.00 mpas_ocean.NeighborhoodDask.track_peakmem_mean('480km', 'numpy')
545k 547k 1.00 mpas_ocean.NeighborhoodDask.track_peakmem_mean('480km', 'time_chunks')
12.7±0.01s 12.5±0.02s 0.99 mpas_ocean.NeighborhoodReduce.time_dataset_reduce('120km', 'mean')
13.3±0.01s 13.2±0.04s 0.99 mpas_ocean.NeighborhoodReduce.time_dataset_reduce('120km', 'median')
227±0.2ms 227±0.8ms 1.00 mpas_ocean.NeighborhoodReduce.time_dataset_reduce('480km', 'mean')
232±1ms 232±2ms 1.00 mpas_ocean.NeighborhoodReduce.time_dataset_reduce('480km', 'median')
1.35±0s 1.34±0s 0.99 mpas_ocean.NeighborhoodReduce.time_neighborhood_reduce('120km', 'mean')
1.54±0.01s 1.54±0.01s 1.00 mpas_ocean.NeighborhoodReduce.time_neighborhood_reduce('120km', 'median')
25.8±0.08ms 25.7±0.09ms 1.00 mpas_ocean.NeighborhoodReduce.time_neighborhood_reduce('480km', 'mean')
27.3±0.2ms 26.9±0.06ms 0.99 mpas_ocean.NeighborhoodReduce.time_neighborhood_reduce('480km', 'median')
233±1ms 228±0.7ms 0.98 mpas_ocean.NeighborhoodReduce.time_reduce('120km', 'median')
453±10μs 454±10μs 1.00 mpas_ocean.NeighborhoodReduce.time_reduce('480km', 'mean')
1.66±0.02ms 1.68±0.02ms 1.01 mpas_ocean.NeighborhoodReduce.time_reduce('480km', 'median')
239k 239k 1.00 mpas_ocean.NeighborhoodReduce.track_peakmem_reduce('120km', 'mean')
245k 245k 1.00 mpas_ocean.NeighborhoodReduce.track_peakmem_reduce('120km', 'median')
19.4k 19.5k 1.00 mpas_ocean.NeighborhoodReduce.track_peakmem_reduce('480km', 'mean')
19.9k 19.9k 1.00 mpas_ocean.NeighborhoodReduce.track_peakmem_reduce('480km', 'median')
406±9μs 400±3μs 0.98 mpas_ocean.PointInPolygon.time_face_search_lonlat('120km')
412±10μs 411±10μs 1.00 mpas_ocean.PointInPolygon.time_face_search_lonlat('480km')
384±20μs 383±10μs 1.00 mpas_ocean.PointInPolygon.time_face_search_xyz('120km')
381±10μs 384±20μs 1.01 mpas_ocean.PointInPolygon.time_face_search_xyz('480km')
134±0.9ms 131±0.9ms 0.98 mpas_ocean.RemapDownsample.time_bilinear_remapping
16.9±0.04ms 17.3±0.2ms 1.02 mpas_ocean.RemapDownsample.time_inverse_distance_weighted_remapping
15.3±0.03ms 15.7±0.2ms 1.02 mpas_ocean.RemapDownsample.time_nearest_neighbor_remapping
1.43±0.01s 1.42±0.01s 0.99 mpas_ocean.RemapUpsample.time_bilinear_remapping
26.6±0.2ms 26.1±0.1ms 0.98 mpas_ocean.RemapUpsample.time_inverse_distance_weighted_remapping
12.4±0.4ms 11.6±0.03ms 0.93 mpas_ocean.RemapUpsample.time_nearest_neighbor_remapping
8.29±0.1ms 8.34±0.2ms 1.01 mpas_ocean.ZonalAverage.time_zonal_average('120km')
5.13±0.07ms 5.14±0.06ms 1.00 mpas_ocean.ZonalAverage.time_zonal_average('480km')
348M 338M 0.97 mpas_ocean.ZonalAveragePeakMem.track_peakmem_zonal_average('120km')
348M 322M 0.93 mpas_ocean.ZonalAveragePeakMem.track_peakmem_zonal_average('480km')
1.0321581221548088 1.0235012720709553 0.99 nogil_scaling.GILScaling.track_gil_scaling
7.09±0.01ms 7.12±0.07ms 1.00 quad_hexagon.QuadHexagon.time_open_dataset
6.05±0.04ms 6.22±0.09ms 1.03 quad_hexagon.QuadHexagon.time_open_grid
408 408 1.00 quad_hexagon.QuadHexagon.track_nbytes_open_dataset
392 392 1.00 quad_hexagon.QuadHexagon.track_nbytes_open_grid
73.2k 73.6k 1.01 quad_hexagon.QuadHexagon.track_peakmem_open_dataset
72.4k 72.7k 1.01 quad_hexagon.QuadHexagon.track_peakmem_open_grid

Benchmarks that have got worse:

Change Before [e01d71e] After [c2b3706] Ratio Benchmark (Parameter)
+ 692±9μs 780±10μs 1.13 mpas_ocean.NeighborhoodDask.time_mean('480km', 'numpy')
+ 685k 757k 1.1 mpas_ocean.NeighborhoodDask.track_peakmem_mean('480km', 'grid_chunks')

@cmdupuis3
cmdupuis3 marked this pull request as ready for review September 18, 2026 21:50

@Sevans711 Sevans711 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The fix looks clean, even though it might look like there are some "object-oriented shenanigans" happening here with making _LazyKernel class a descriptor (i.e., it has a __get__ method), it makes sense to me and the comments communicate the intent clearly, so it shouldn't be too difficult to maintain in the future if needed. I think this is a good way to proceed.

Noticed the small slowdowns (~20%) in some neighborhood benchmarks, but I think that is acceptable if this PR is solving the race condition and/or lack of caching (which slows down real workflows the first time it hits but isn't necessarily appearing in the benchmark suite).

The only issue is, I still have the crash from the race condition (#1757)! Copy-pasting those lines into terminal while on this PR's branch, it still crashes locally for me within a few runs.

If you want to merge this as closing #1769 and related to 1757 but not closing 1757, I would be okay to approve. But for closing 1757 something else needs to be fixed. I'm not sure what is missing (race conditions are tricky to debug…) but I wonder if it might be as simple as moving the threading.Lock() mechanism (not sure where it would need to move to, though…)?

@cmdupuis3

Copy link
Copy Markdown
Collaborator Author

pre-commit.ci autofix

@Sevans711 Sevans711 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks like the race condition is now fixed as well! Running the code from #1757, I now see no crashes after 20 runs, whereas before it was usually crashing within a few runs at most, and never took more than 10 runs before the race condition caused a crash.

Skimmed through the code changes since last time I viewed and they look reasonable.

My only concern now is that some of the 480km Neighborhood benchmarks show slowdowns. Though, the 120km ones show speedups…? This makes me wonder if maybe the ways this PR affects when numba compilation happens here is causing the ASV benchmarking suite to not warm up properly in this case. Are these slowdowns real, or just because they are including the numba compilation times? Aside from that, I would be happy to approve!

@cmdupuis3

Copy link
Copy Markdown
Collaborator Author

@Sevans711 Good catch, there was some numba compilation getting into the timings on smaller grids. Otherwise, there is a small overall speedup on the neighborhood reductions. Some of the other benchmark differences are probably noise (like the open_grid, since I didn't touch that)

@Sevans711 Sevans711 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Approving, now that the benchmarks look better!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

benchmarking Related to benchmarks, memory usage, and/or time profiling bug Something isn't working run-benchmark Run ASV benchmark workflow scalability Related to scalability & performance efforts

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Neighborhood filter kernels don't actually cache Race condition crashes pytest suite locally sometimes

2 participants